use proc_macro::TokenStream;
use quote::quote;
use syn::{parse_macro_input, spanned::Spanned, DeriveInput, Fields, Type};

#[proc_macro_derive(VariantMaybeDeserOwned)]
pub fn derive_variant_maybe_deser_owned(item: TokenStream) -> TokenStream {
    let mut inp = parse_macro_input!(item as DeriveInput);
    // impl (generic) VariantMaybeDeserOwned for <name> (generic) (generic_where) {
    //     fn variant_maybe_deser_owned(&self) -> {
    //         match self {
    //             <name>::<variant1>((..)) => ...,
    //         }
    //     }
    // }

    let syn::Data::Enum(ref mut data) = inp.data else {
        return syn::parse::Error::new(
            inp.span(),
            "VariantMaybeDeserOwned should only be applied to enums!",
        )
        .into_compile_error()
        .into();
    };
    let name = inp.ident;

    let before_our_where = if let Some(ref mut their_where) = inp.generics.where_clause {
        their_where.predicates.pop_punct();
        quote!(,)
    } else {
        quote!(where)
    };
    let (impl_generic, type_generic, their_where) = inp.generics.split_for_impl();
    let type_param_names = inp.generics.type_params().map(|x| &x.ident);
    data.variants
        .iter_mut()
        .flat_map(|variant| &mut variant.fields)
        .for_each(|field| {
            if let Type::Reference(ref mut r) = field.ty {
                r.lifetime = None;
            }
        });
    let variant_match_pats = data.variants.iter().map(|variant| {
        let variant_name = &variant.ident;
        match variant.fields {
            Fields::Named(..) => quote!(#variant_name { .. }),
            Fields::Unnamed(..) => quote!(#variant_name (..)),
            Fields::Unit => quote!(#variant_name),
        }
    });
    let variant_checks = data.variants.iter().map(|variant| {
        let inner_types = variant.fields.iter().map(|field| &field.ty);
        quote! {
            true #(&& <#inner_types as MaybeDeserOwned>::IS_DESER_OWNED)*
        }
    });
    quote! {
        unsafe impl #impl_generic VariantMaybeDeserOwned for #name #type_generic
            #their_where #before_our_where #(#type_param_names: MaybeDeserOwned),*
        {
            fn variant_is_deser_owned(&self) -> bool {
                match self {
                    #(#name::#variant_match_pats => #variant_checks,)*
                }
            }
        }
    }
    .into()
}

#[proc_macro_derive(MaybeDeserOwned)]
pub fn derive_maybe_deser_owned(item: TokenStream) -> TokenStream {
    let mut inp = parse_macro_input!(item as DeriveInput);

    // impl (generic) MaybeDeserOwned for <name> (generic) (generic_where) {
    //     const IS_DESER_OWNED: bool = ...;
    // }

    let name = inp.ident;

    let before_our_where = if let Some(ref mut their_where) = inp.generics.where_clause {
        their_where.predicates.pop_punct();
        quote!(,)
    } else {
        quote!(where)
    };
    let (impl_generic, type_generic, their_where) = inp.generics.split_for_impl();
    let type_param_names = inp.generics.type_params().map(|x| &x.ident);
    let mut out = quote! {
        unsafe impl #impl_generic MaybeDeserOwned for #name #type_generic
            #their_where #before_our_where #(#type_param_names: MaybeDeserOwned),*
    };

    macro_rules! impl_body {
        ($iter: ident, $inp: ident, $out: ident) => {
            let inner_types = $iter.map(|field| {
                if let Type::Reference(ref mut r) = field.ty {
                    r.lifetime = None;
                }
                &field.ty
            });
            let type_param_names = $inp
                .generics
                .type_params()
                .map(|x| &x.ident);
            $out.extend(quote! {
                {
                    const IS_DESER_OWNED: bool = {
                        let first = true #(&& <#type_param_names as MaybeDeserOwned>::IS_DESER_OWNED)*;
                        let second = true #(&& <#inner_types as MaybeDeserOwned>::IS_DESER_OWNED)*;
                        first && second
                    };
                }
            });
        };
    }
    match inp.data {
        syn::Data::Struct(ref mut st) => {
            let iter = st.fields.iter_mut();
            impl_body!(iter, inp, out);
        }
        syn::Data::Enum(ref mut en) => {
            let iter = en
                .variants
                .iter_mut()
                .flat_map(|variant| &mut variant.fields);
            impl_body!(iter, inp, out);
        }
        syn::Data::Union(ref mut un) => {
            let iter = un.fields.named.iter_mut();
            impl_body!(iter, inp, out);
        }
    }
    out.into()
}
